package com.xiashitech.agent.instrumentation.httprequest.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.constants.Instrumentation;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSRequestCallDepth;
import com.xiashitech.agent.instrumentation.httprequest.XSHttpServletRequestWrapper;
import com.xiashitech.agent.instrumentation.httprequest.XSHttpServletResponseWrapper;
import com.xiashitech.agent.instrumentation.httprequest.utils.SpecialUtil;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import com.xiashitech.agent.utils.SessionUtil;
import com.xiashitech.agent.utils.TypeUtils;
import com.xiashitech.interfaces.agent.dto.AgentConfigDTO;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Optional;

import static com.xiashitech.agent.instrumentation.httprequest.filter.XSRequestSingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSRequestInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return hasSuperType(named("javax.servlet.Filter"));
    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        typeTransformer.applyAdviceToMethod(
                isMethod().and(namedOneOf("doFilter"))
                        .and(takesArgument(0, named("javax.servlet.ServletRequest")))
                        .and(takesArgument(1, named("javax.servlet.ServletResponse")))
                        .and(takesArgument(2, named("javax.servlet.FilterChain")))
                        .and(isPublic()),
                XSRequestInstrumentation.class.getName() + "$FilterAdvice");
    }

    public static class FilterAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
                @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("user") UserRequest userRequest) {
            DurationPrinter durationPrinter = new DurationPrinter("XSRequestInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSRequestCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            try {
                if(!(request instanceof HttpServletRequest && response instanceof HttpServletResponse)) {
                    return;
                }
                HttpServletRequest httpServletRequest = (HttpServletRequest)request;
                // 对于静态资源，如 css、image、js 等不做处理
                String acceptType = httpServletRequest.getHeader("Accept");
                if(TypeUtils.isSkipAcceptType(acceptType, httpServletRequest.getRequestURI())) {
                    return;
                }
                //String loginPathKey = AgentConfigDTO.getConfig().getLoginPath().get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                Boolean httpWrapper = null;
                if(AgentConfigDTO.getConfig().getHttpWrapper() != null) {
                    httpWrapper = AgentConfigDTO.getConfig().getHttpWrapper().get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                }

                //只要请求头里content-type是json的就将其包装为XsRequest
                if((httpWrapper == null || httpWrapper) && (httpServletRequest.getContentType() != null && httpServletRequest.getContentType().toLowerCase().contains("json"))) {
                    request = new XSHttpServletRequestWrapper(httpServletRequest);
                }
                if((httpWrapper == null || httpWrapper) && !SpecialUtil.specialRequestForResponseWrapper(httpServletRequest.getRequestURI())) {
                    response = new XSHttpServletResponseWrapper((HttpServletResponse)response);
                }
                Optional<UserRequest> userRequestOptional = Java8Bridge.executeRuleForUserInfoMatchRule((HttpServletRequest)request);
                if(userRequestOptional.isPresent()) {
                    userRequest = userRequestOptional.get();
                } else {
                    userRequest = UserRequest.builder().build();
                }
                userRequest.setRequest(request);
                Context parentContext = currentContext();
                if(Java8Bridge.currentForBaggage().getEntryValue("xsFlowStart") == null) {
                    Baggage xsFlowStartBaggage = Java8Bridge.currentForBaggage().toBuilder()
                            .put("xsFlowStart", "true")
                            .put("xsAuthSystemId", StartConfig.getParamValueFromKey(StartConfig.systemId))
                            .build();
                    parentContext = parentContext.with(xsFlowStartBaggage);
                }
                if (instrumenter().shouldStart(parentContext,userRequest)) {
                    parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                    context = instrumenter().start(parentContext, userRequest);
                    String accessToken = AgentConfigDTO.getConfig().getAccessToken().get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                    String token = "";
                    if(accessToken != null && !"".equals(accessToken)) {
                        token = httpServletRequest.getHeader(accessToken);
                        if(token != null && !"".equals(token) && token.startsWith(Instrumentation.AUTHORIZATION_BEARER)) {
                            int index = token.indexOf(Instrumentation.AUTHORIZATION_BEARER + " ");
                            if(index >= 0) {
                                token = token.substring(index + 7).trim();
                            }
                        } else if(token == null) {
                            token = SessionUtil.getCookieValueByKey(httpServletRequest,accessToken);
                        }
                    }
                    // userRequest.setLoginName(token);
                    context = BaggageUtil.syncBaggage()
                            .put("xs.userId",userRequest.getUserid() == null ? "" : userRequest.getUserid())
                            .put("xs.loginName",userRequest.getLoginName() == null ? "" : userRequest.getLoginName())
                            .put("xs.sessionId",userRequest.getSessionId() == null ? "" : userRequest.getSessionId())
                            .put("xs.requestDate",userRequest.getRequestDate() == null ? "" : userRequest.getRequestDate().toString())
                            .put("xs.sessionUpdate",userRequest.getSessionUpdate() == null ? "false" : userRequest.getSessionUpdate().toString())
                            .put("xs.logintoken",token == null || "".equals(token) ? "" : token)
                            .build().storeInContext(context);
                    scope = context.makeCurrent();
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("FilterAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
                @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Thrown Throwable throwable,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("user") UserRequest userRequest) {
            DurationPrinter durationPrinter = new DurationPrinter("XSRequestInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            try {
                if(!(request instanceof HttpServletRequest && response instanceof HttpServletResponse)) {
                    return;
                }
                if(response instanceof XSHttpServletResponseWrapper) {
                    ((XSHttpServletResponseWrapper)response).transferDataToResponse();
                }
                if (scope == null) {
                    return;
                }
                if(userRequest.getUserid() == null || "".equals(userRequest.getUserid())) {
                    Optional<UserRequest> userRequestOptional = Java8Bridge.executeRuleForUserInfoMatchRule((HttpServletRequest)request);
                    if(userRequestOptional.isPresent() && userRequestOptional.get().getUserid() != null && !"".equals(userRequestOptional.get().getUserid())) {
                        userRequest = userRequestOptional.get();
                        userRequest.setSessionUpdate(true);
                    }
                }
                scope.close();
                instrumenter().end(context,userRequest,response,throwable);
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("FilterAdvice e");
        }
    }
}
